package assignment02;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

import nlp.classify.LabeledInstance;
import nlp.classify.ProbabilisticClassifier;
import nlp.classify.ProbabilisticClassifierFactory;
import nlp.util.Counter;
import nlp.util.CounterMap;

public class NaiveBayesClassifier implements ProbabilisticClassifier<String, String>{

	public static class Factory implements ProbabilisticClassifierFactory<String, String>{
		private boolean usingSmoothing;
		private double alpha;
		private int order;
		private double smoothingFactor;
		public Factory(boolean usingSmoothing, double alpha, int order, double smoothingFactor) {
			ngramMaps = new HashMap<String, CounterMap<String,Character>>();
			suffixMap = new CounterMap<String, String>();
			numOfWordsMap = new CounterMap<String, Integer>();
			lengthMap = new CounterMap<String, Integer>();
			wordsMap = new CounterMap<String, String>();
			vocabularyMap = new CounterMap<String, Character>();
			categoryMap = new Counter<String>();
			this.usingSmoothing = usingSmoothing;
			this.alpha = alpha;
			this.order = order;
			this.smoothingFactor = smoothingFactor;
		}
		
		private static final char START = '\uE000';
		private static final double INCREMENT = 1.0;
		private static final String SPLIT_PATTERN = "[ \t\n\r-]";
		
		private Map<String, CounterMap<String, Character>> ngramMaps;
		private static final String NOSUFFIX = "*NOSUFFIX*";
		private CounterMap<String, String> suffixMap;
		private CounterMap<String, Integer> numOfWordsMap;
		private CounterMap<String, Integer> lengthMap;
		private CounterMap<String, String> wordsMap;
		private CounterMap<String, Character> vocabularyMap;
		private Counter<String> categoryMap;
		
		public ProbabilisticClassifier<String, String> trainClassifier(
				List<LabeledInstance<String, String>> trainingData) {
			// TODO Auto-generated method stub
			for (LabeledInstance<String, String> datum : trainingData) {
				String name = datum.getInput().toString();
				String label = datum.getLabel().toString();
				ngramFeature(label, name, order);
				suffixFeature(label, name);
				numOfWordsFeature(label, name);
				lengthFeature(label, name);
				wordsFeature(label, name);
//				categoryMap.incrementCount(label, INCREMENT);
			}
			// count
			Counter<String> counter = new Counter<String>();
			for (LabeledInstance<String, String> datum : trainingData) {
				counter.incrementCount(datum.getLabel(), INCREMENT);
			}
			for (String label : counter.keySet()) {
				System.out.println(label + " & " + counter.getCount(label));
			}
			categoryProbability();

//			System.out.println("training done");
			return new NaiveBayesClassifier(ngramMaps, suffixMap, numOfWordsMap, lengthMap, wordsMap, vocabularyMap, categoryMap, usingSmoothing, alpha, smoothingFactor, order);
		}
		
		
		private void ngramFeature(String label, String name, int order) {
			CounterMap<String, Character> ngramMap = ensureMap(label);
			Counter<Character> vocabularyCounter = vocabularyMap.getCounter(label);
			String prefix = "";
			for (int i = 0; i < order; i++) {
				prefix += START;
			}
				prefix += name;
			for (int i = 0; i < prefix.length() - order; i++) {
				String history = prefix.substring(i, i+order-1);
				char ch = prefix.charAt(i+order);
				Counter<Character> chCounter = ngramMap.getCounter(history);
				vocabularyCounter.setCount(ch, INCREMENT);
				chCounter.incrementCount(ch, INCREMENT);
			}
		}
		
		private void suffixFeature(String label, String name) {
			Counter<String> suffixCounter = suffixMap.getCounter(label);
			String[] words = name.split(SPLIT_PATTERN);
			if (words.length > 1) {
				String suffix = words[words.length - 1];
				suffixCounter.incrementCount(suffix, INCREMENT);
			} else {
				suffixCounter.incrementCount(NOSUFFIX, INCREMENT);
			}
		}
		
		private void numOfWordsFeature(String label, String name) {
			Counter<Integer> wordsCounter = numOfWordsMap.getCounter(label);
			String[] words = name.split(SPLIT_PATTERN);
			int numOfWords = words.length;
			wordsCounter.incrementCount(numOfWords, INCREMENT);
		}
		
		private void lengthFeature(String label, String name) {
			Counter<Integer> lengthCounter = lengthMap.getCounter(label);
			int length = name.length();
			lengthCounter.incrementCount(length, INCREMENT);
		}
		
		private void wordsFeature(String label, String name) {
			Counter<String> counter = wordsMap.getCounter(label);
			String[] words = name.split("[ \t\r\n\f]");
			for (int i = 0; i < words.length; i++) {
				counter.incrementCount(words[i], INCREMENT);
			}
		}
		
		private void categoryProbability() {
			for (String label : ngramMaps.keySet()) {
				CounterMap<String, Character> counterMap = ngramMaps.get(label);
				double total = counterMap.totalCount();
				categoryMap.incrementCount(label, total);
			}
			for (String label : suffixMap.keySet()) {
				Counter<String> counter = suffixMap.getCounter(label);
				double total = counter.totalCount();
				categoryMap.incrementCount(label, total);
			}
			for (String label : numOfWordsMap.keySet()) {
				Counter<Integer> counter = numOfWordsMap.getCounter(label);
				double total = counter.totalCount();
				categoryMap.incrementCount(label, total);
			}
			for (String label : lengthMap.keySet()) {
				Counter<Integer> counter = lengthMap.getCounter(label);
				double total = counter.totalCount();
				categoryMap.incrementCount(label, total);
			}
			for (String label : wordsMap.keySet()) {
				Counter<String> counter = wordsMap.getCounter(label);
				double total = counter.totalCount();
				categoryMap.incrementCount(label, total);
			}
			categoryMap.normalize();
		}
		
		/** 
		 * eusure a CounterMap object
		 * @param label
		 * @return
		 */
		private CounterMap<String, Character> ensureMap(String label) {
			if (!ngramMaps.containsKey(label)) {
				CounterMap<String, Character> counterMap = new CounterMap<String, Character>();
				ngramMaps.put(label, counterMap);
			} 

			return ngramMaps.get(label);
		}
		
	}

	
	private static final char START = '\uE000';
	private static final String NOSUFFIX = "*NOSUFFIX*";
	private static final String SPLIT_PATTERN = "[ \t\n\r-]";
	private double SMOOTHINGFACTOR = 1.0;

	private Map<String, CounterMap<String, Character>> ngramMaps;
	private CounterMap<String, String> suffixMap;
	private CounterMap<String, Integer> numOfWordsMap;
	private CounterMap<String, Integer> lengthMap;
	private CounterMap<String, String> wordsMap;
	private CounterMap<String, Character> vocabularyMap;
	private Counter<String> categoryMap;
	private Counter<String> vocabularyCounter;
	
	private boolean usingSmoothing = false;
	private double alpha = 0.0;
	private int order; //= 3;
	
	public NaiveBayesClassifier(Map<String, CounterMap<String, Character>> ngramMaps,
			CounterMap<String, String> suffixMap,
			CounterMap<String, Integer> numOfWordsMap,
			CounterMap<String, Integer> lengthMap,
			CounterMap<String, String> wordsMap,
			CounterMap<String, Character> vocabularyMap,
			Counter<String> categoryMap,
			boolean usingSmoothing,
			double alpha,
			double smoothingFactor,
			int order) {
		this.ngramMaps = ngramMaps;
		this.suffixMap = suffixMap;
		this.numOfWordsMap = numOfWordsMap;
		this.lengthMap = lengthMap;
		this.wordsMap = wordsMap;
		this.vocabularyMap = vocabularyMap;
		this.categoryMap = categoryMap;
		
		initVocabularyCounter();
		this.usingSmoothing = usingSmoothing;
		this.alpha = alpha;
		this.SMOOTHINGFACTOR = smoothingFactor;
		this.order = order;
	}
	
	void initVocabularyCounter() {
		vocabularyCounter = new Counter<String>();
		for (String label : vocabularyMap.keySet()) {
			Counter<Character> counter = vocabularyMap.getCounter(label);
			if (counter.size() == 0) {
				System.out.println("shit");
			}
			vocabularyCounter.setCount(label, counter.size());
		}
	}
	
	
	public String getLabel(String instance) {
		// TODO Auto-generated method stub
		return getProbabilities(instance).argMax();
	}

	public Counter<String> getProbabilities(String instance) {
		// TODO Auto-generated method stub
		Counter<String> counter = new Counter<String>();
		for (String label : categoryMap.keySet()) {
			double p = getINSTANCELogProbability(label, instance);
			double count = p + alpha * getLABELLogProbability(label);
			double probability = Math.exp(count);
			counter.setCount(label, probability);
//			System.out.println("probability is " + count);
		}
		double totalCount = counter.totalCount();
		for (String label : categoryMap.keySet()) {
			double probability = counter.getCount(label) / totalCount;
			counter.setCount(label, probability);
//			System.out.println("probability is " + count);
		}
//		System.out.println();
		return counter;
	}
	
	private double getINSTANCELogProbability(String label, String name) {
		double logProbability = 0.0;
		logProbability += getNGRAMLogProbability(label, name);
		logProbability += getSUFFIXLogProbability(label, name);
		logProbability += getNUMOFWORDSLogProbability(label, name);
		logProbability += getLENGTHLogProbability(label, name);
		logProbability += getWORDSLogProbability(label, name);
		return logProbability;
	}
	
	private double getNGRAMLogProbability(String label, String name) {
		CounterMap<String, Character> ngramMap = ngramMaps.get(label);
		String prefix = "";
		for (int i = 0; i < order; i++) {
			prefix += START;
		}
		prefix += name;
		double vocabulary = vocabularyCounter.getCount(label);
		double logProbability = 0.0;
		for (int i = 0; i < prefix.length() - order; i++) {
			String history = prefix.substring(i, i+order-1);
			char ch = prefix.charAt(i+order);
			Counter<Character> chCounter = ngramMap.getCounter(history);
			double uniCount = chCounter.getCount(ch);
			double totalCount = chCounter.totalCount();
			if (usingSmoothing) {
				uniCount += SMOOTHINGFACTOR;
				totalCount += SMOOTHINGFACTOR * vocabulary;
			}
			logProbability += Math.log(uniCount / totalCount);
			if (Double.isInfinite(logProbability)) {
				System.out.println("getNGRAMLogProbability: " + logProbability);
			}
		}
		return logProbability;
	}
	
	private double getSUFFIXLogProbability(String label, String name) {
		Counter<String> suffixCounter = suffixMap.getCounter(label);
		String[] words = name.split(SPLIT_PATTERN);
		double count = 0.0;
		double totalCount = suffixCounter.totalCount();

		if (words.length == 1) {
			count = suffixCounter.getCount(NOSUFFIX);
		} else {
			count = suffixCounter.getCount(words[words.length - 1]);
		}
		if (usingSmoothing) {
			count += SMOOTHINGFACTOR;
			totalCount += SMOOTHINGFACTOR * suffixCounter.size();
		}
		double logProbability = Math.log(count / totalCount);
		if (Double.isInfinite(logProbability)) {
			System.out.println("getSUFFIXLogProbability: " + logProbability);
		}
		return logProbability;
	}
	
	private double getNUMOFWORDSLogProbability(String label, String name) {
		Counter<Integer> counter = numOfWordsMap.getCounter(label);
		String[] words = name.split(SPLIT_PATTERN);
		int numWords = words.length;
		double count = counter.getCount(numWords);
		double totalCount = counter.totalCount();
		if (usingSmoothing) {
			count += SMOOTHINGFACTOR;
			totalCount += SMOOTHINGFACTOR * counter.size();
		}
		double logProbability = Math.log(count / totalCount);
		if (Double.isInfinite(logProbability)) {
			System.out.println("getNUMOFWORDSLogProbability: " + logProbability);
		}
		return logProbability;
	}
	
	private double getLENGTHLogProbability(String label, String name) {
		Counter<Integer> counter = lengthMap.getCounter(label);
		int length = name.length();
		double count = counter.getCount(length);
		double totalCount = counter.totalCount();
		if (usingSmoothing) {
			count += SMOOTHINGFACTOR;
			totalCount += SMOOTHINGFACTOR * counter.size();
		}
		double logProbability = Math.log(count / totalCount);
		if (Double.isInfinite(logProbability)) {
			System.out.println("getLENGTHLogProbability: " + logProbability);
		}
		return logProbability;
	}
	
	private double getWORDSLogProbability(String label, String name) {
		Counter<String> counter = wordsMap.getCounter(label);
		String[] words = name.split(SPLIT_PATTERN);
		double logProbability = 0.0;
		double totalCount = counter.totalCount();
		for (int i = 0; i < words.length; i++) {
			String word = words[i];
			double count = counter.getCount(word);
			if (usingSmoothing) {
				count += SMOOTHINGFACTOR;
				int size = counter.size();
				double p = Math.log(count / (totalCount + SMOOTHINGFACTOR * size));
				logProbability += p;
			}
		}
		if (Double.isInfinite(logProbability)) {
			System.out.println("getWORDSLogProbability: " + logProbability);
		}
		return logProbability;
	}
	
	private double getLABELLogProbability(String label) {
		double count = categoryMap.getCount(label);
		double totalCount = categoryMap.totalCount();
		double logProbability = Math.log(count / totalCount);
//		System.out.println("labal probability is " + logProbability);
		return logProbability;
	}
}
